Align interleaved source locators with Parquet FILE by VibhuJawa · Pull Request #2284 · NVIDIA-NeMo/Curator

VibhuJawa · 2026-08-01T06:08:31Z

Summary

Align interleaved source locators with the Parquet FILE logical-type shape, simplify materialization, and improve actual Parquet reader/writer throughput.

Previous field	New location
`source_ref.path`	`source_ref.uri`
`source_ref.byte_offset`	`source_ref.offset`
`source_ref.byte_size`	`source_ref.size`
`source_ref.member`	adjacent `source_member` column
`source_ref.frame_index`	adjacent `source_frame_index` column

source_ref has all six optional FILE children from the specification: uri, offset, size, content_type, checksum, and inline.

What changed

WebDataset emits a typed FILE-compatible struct, preserving canonical Arrow child order.
FILE_REFERENCE_TYPE lives in shared nemo_curator.utils.storage_utils for reuse by text, image, video, and interleaved stages.
Existing JSON-string locators are migrated once at the shared schema-alignment boundary.
Legacy Nemotron {page, bbox} provenance migrates to typed adjacent page_number and source_bbox columns instead of being discarded.
uri + size now correctly resolves [0, size) when offset is absent.
Whole-object and ranged references share one deduplicated fs.cat_ranges() path, with one call per filesystem backend.
Parquet output is opened through the configured fsspec filesystem, preserving remote protocols and storage_options.
Arrow inputs bypass redundant Arrow → pandas → Arrow conversion when materialization/error filtering is unnecessary.
Dictionary encoding defaults to low-cardinality interleaved fields; callers can still override use_dictionary.

Archive-member and TIFF-frame metadata remain adjacent because FILE is a closed group. source_member is only needed for tar extraction when a usable FILE byte range is unavailable; ranged and whole-object FILE reads do not depend on it.

Compatibility and scope

The change keeps binary_content unchanged. The inline FILE child is present in the canonical schema, but moving existing payload storage or materialization to inline is outside this PR.

PyArrow 19 can write the exact physical group but cannot emit the new LogicalType.FILE footer annotation. The FILE specification merged on 2026-07-26, and the cited parquet-java and arrow-rs implementations are still open/unreleased. This PR therefore avoids a custom Thrift footer rewriter and can adopt native Arrow support when it lands.

Performance validation

Compared latest upstream main (5732455) with this branch on the same CPU host. These are real InterleavedParquetWriterStage and InterleavedParquetReaderStage runs over 200,000 typed Arrow rows; values are medians of seven warm runs.

Metric	Upstream main	This PR	Change
Parquet write throughput	4.78M rows/s	6.24M rows/s	+30.7%
Parquet read throughput	3.43M rows/s	4.03M rows/s	+17.6%
Parquet file size	3,128,791 bytes	1,592,172 bytes	−49.1%

Remote HTTP materialization used a pinned, range-capable 466 KB object and verified every returned payload:

32 rows / 8 unique byte ranges: combined alternating-run median was 0.109 s here versus 0.115 s upstream (network-sensitive, no regression).
32 duplicate whole-object references: combined median was 0.174 s here versus 0.273 s upstream; individual runs were network-noisy, but the unified path showed no consistent slowdown.

Validation

Ruff passes across nemo_curator and interleaved tests.
208 top-level interleaved CPU tests pass, covering locator and PDF-provenance migration, FILE semantics, size-only ranges, materialization, fsspec output, Parquet/WebDataset/Lance readers and writers, and round trips.
Latest fork main, upstream main, and the PR base all resolve to 5732455; the branch is based directly on current main.

Diff

19 files changed
356 additions / 359 deletions
net reduction of 3 lines

copy-pr-bot · 2026-08-01T06:08:34Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-08-01T06:34:38Z

Greptile Summary

The PR replaces JSON-string interleaved source locators with a Parquet FILE-compatible struct while moving archive, frame, and PDF provenance into adjacent typed columns. It also consolidates legacy migration at schema alignment and streamlines materialization and Parquet I/O.

Emits canonical FILE-shaped source references from WebDataset ingestion.
Migrates legacy locator and Nemotron provenance fields during shared schema alignment.
Batches and deduplicates whole-object and ranged materialization requests.
Writes Arrow tables directly through the configured fsspec filesystem.

Confidence Score: 5/5

The PR appears safe to merge because the previously reported compatibility and provenance failures are resolved at the current head.

No blocking failure remains.

Important Files Changed

Filename	Overview
nemo_curator/stages/interleaved/utils/schema.py	Centralizes migration of legacy JSON locators and PDF provenance before schema reconciliation; the previously reported compatibility and provenance failures are addressed.
nemo_curator/tasks/interleaved.py	Defines the canonical FILE-compatible source_ref shape and adjacent archive/frame metadata columns, with current callers updated to the replacement API.
nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py	Preserves Nemotron page and bounding-box provenance in typed adjacent columns rather than overloading source_ref.
nemo_curator/stages/interleaved/utils/materialization.py	Updates materialization to consume FILE locators and adjacent metadata while sharing deduplicated filesystem range requests.
nemo_curator/stages/interleaved/io/writers/base.py	Adds a direct Arrow-table preparation path while retaining materialization-error handling when conversion is necessary.
nemo_curator/stages/interleaved/io/writers/tabular.py	Writes Parquet through the configured fsspec filesystem and applies default dictionary encoding to selected low-cardinality fields.
nemo_curator/stages/interleaved/io/readers/webdataset.py	Emits typed FILE references and stores tar-member and TIFF-frame metadata in adjacent canonical columns.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Legacy JSON source_ref] --> B[align_interleaved_table]
  C[Typed WebDataset locator] --> B
  B --> D[FILE-compatible source_ref]
  B --> E[source_member]
  B --> F[source_frame_index]
  B --> G[page_number and source_bbox]
  D --> H[Materialization]
  E --> H
  F --> H
  D --> I[Parquet writer via fsspec]
  E --> I
  F --> I
  G --> I

_{Reviews (9): Last reviewed commit: "refactor(schema): share Parquet FILE typ..." | Re-trigger Greptile}

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

VibhuJawa · 2026-08-01T20:43:22Z

/okay to test 7a5a8e9

VibhuJawa · 2026-08-02T01:09:28Z

/claude review

Copilot

Pull request overview

This PR migrates interleaved provenance from a JSON-string source_ref locator to a Parquet FILE-shaped source_ref struct (with adjacent source_member / source_frame_index), and updates schema alignment, materialization, and I/O stages to improve correctness and throughput while preserving backward compatibility via a one-time migration at the schema boundary.

Changes:

Introduces a canonical Arrow FILE_REFERENCE_TYPE and updates INTERLEAVED_SCHEMA to store source_ref as a FILE-compatible struct plus adjacent metadata columns.
Adds align_interleaved_table() migration/casting logic and updates materialization + readers/writers to operate on typed Arrow tables and fsspec-backed output.
Updates tests and docs to reflect the new locator shape and semantics (including PDF provenance migration to adjacent columns).

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/stages/interleaved/test_schema_utils.py	Adds coverage for legacy JSON `source_ref` migration into typed FILE struct + adjacent columns.
tests/stages/interleaved/test_multimodal_writer.py	Updates writer tests for typed `source_ref`, adjacent member/frame fields, and fsspec-backed Parquet output.
tests/stages/interleaved/test_multimodal_reader.py	Updates reader assertions to validate `source_member` / `source_frame_index` instead of parsing JSON.
tests/stages/interleaved/test_multimodal_core.py	Adjusts core materialization/classification tests to the FILE + adjacent metadata model and unified range-read path.
tests/stages/interleaved/test_materialization.py	Extends edge-case coverage (e.g., size-only reads; empty-range acceptance) under new range-read semantics.
tests/stages/interleaved/test_interleaved_task.py	Verifies canonical FILE child ordering and validates range constraints in `build_source_ref`.
tests/stages/interleaved/test_base_writer.py	Updates base-writer tests to `_write_table` and Arrow-table alignment outputs.
tests/stages/interleaved/pdf/nemotron_parse/test_utils.py	Adds assertions for migrated PDF provenance fields (`page_number`, `source_bbox`) with `source_ref=None`.
tests/stages/interleaved/conftest.py	Updates fixtures/helpers to construct typed FILE `source_ref` and adjacent metadata columns.
nemo_curator/tasks/interleaved.py	Defines `FILE_REFERENCE_TYPE`, updates `INTERLEAVED_SCHEMA`, and replaces JSON helpers with FILE struct helpers/expansion.
nemo_curator/stages/interleaved/utils/schema.py	Adds `align_interleaved_table()` to migrate/cast legacy locator shapes to canonical FILE struct and append adjacent columns.
nemo_curator/stages/interleaved/utils/materialization.py	Consolidates whole-file and ranged reads into deduped `fs.cat_ranges()` flow; updates classification inputs/columns.
nemo_curator/stages/interleaved/README.md	Documents new FILE-shaped `source_ref` and adjacent metadata strategy; updates materialization strategy description.
nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py	Migrates Nemotron provenance from embedded JSON `source_ref` to adjacent `page_number` / `source_bbox`.
nemo_curator/stages/interleaved/io/writers/webdataset.py	Moves writer implementation to `_write_table` and converts Arrow → pandas only at the boundary.
nemo_curator/stages/interleaved/io/writers/tabular.py	Switches Parquet writing to `pq.write_table` through the configured fsspec filesystem and sets dictionary defaults.
nemo_curator/stages/interleaved/io/writers/base.py	Refactors base writer to prepare/align Arrow tables directly and fast-path Arrow inputs when materialization is unnecessary.
nemo_curator/stages/interleaved/io/readers/webdataset.py	Emits typed FILE-compatible `source_ref` scalars and populates adjacent `source_member` / `source_frame_index`.

+        refs = [json.loads(value) if value else {} for value in table.column(source_ref_index).to_pylist()]
+        for ref in refs:
+            if (path := ref.get("path")) is not None:
+                ref.update(
+                    uri=str(path),
+                    offset=int(ref["byte_offset"]) if ref.get("byte_offset") is not None else None,
+                    size=int(ref["byte_size"]) if ref.get("byte_size") is not None else None,
+                )
+        table = table.set_column(
+            source_ref_index,
+            "source_ref",
+            pa.array((ref if ref.get("path") is not None else None for ref in refs), type=FILE_REFERENCE_TYPE),
+        )


claude · 2026-08-02T01:17:21Z

+        if (
+            ctx.member_info
+            and content_key in ctx.member_info
+            and not ctx.tar_path.lower().split("?", 1)[0].endswith((".tar.gz", ".tgz"))


This .tar.gz/.tgz guard is a subtle correctness fix — info.offset_data/info.size are positions in the decompressed stream, so emitting them as a FILE byte range would make materialization cat_ranges() read garbage from the compressed file. Good catch, but it's new behavior with no direct test. Consider adding a case that reads image members from a gzipped tar and asserts source_ref carries no offset/size (falls back to tar-extract), so a future refactor doesn't silently reintroduce the range read.

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

VibhuJawa force-pushed the codex/interleaved-file branch from 3276dbc to 7ec6416 Compare August 1, 2026 06:26

VibhuJawa changed the title ~~Align interleaved source refs with Parquet FILE~~ Align interleaved source locators with Parquet FILE Aug 1, 2026

VibhuJawa requested a review from suiyoubi August 1, 2026 06:31

VibhuJawa marked this pull request as ready for review August 1, 2026 06:32

VibhuJawa requested a review from a team as a code owner August 1, 2026 06:32

greptile-apps Bot reviewed Aug 1, 2026

View reviewed changes

Comment thread nemo_curator/tasks/interleaved.py

Comment thread nemo_curator/tasks/interleaved.py

VibhuJawa added 2 commits August 1, 2026 07:01

Align interleaved source refs with Parquet FILE

e1f6322

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

Handle legacy interleaved source refs

a58f01d

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

VibhuJawa force-pushed the codex/interleaved-file branch from 528e8a5 to a58f01d Compare August 1, 2026 07:02

VibhuJawa added 2 commits August 1, 2026 17:07

Prioritize FILE byte ranges during materialization

5fa21c8

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

Fix FILE I/O correctness and performance

8fad1aa

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

greptile-apps Bot reviewed Aug 1, 2026

View reviewed changes

Comment thread nemo_curator/stages/interleaved/utils/schema.py Outdated

VibhuJawa force-pushed the codex/interleaved-file branch from d86a4a6 to 8fad1aa Compare August 1, 2026 20:01

fix(interleaved): preserve PDF provenance

7a5a8e9

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

copy-pr-bot Bot temporarily deployed to public August 1, 2026 20:43 Inactive

copy-pr-bot Bot temporarily deployed to test August 1, 2026 20:43 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci August 1, 2026 20:43 Inactive

copy-pr-bot Bot temporarily deployed to public August 1, 2026 20:48 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci August 1, 2026 20:54 Inactive

VibhuJawa requested a review from Copilot August 2, 2026 01:09

Copilot started reviewing on behalf of VibhuJawa August 2, 2026 01:09 View session

Copilot AI reviewed Aug 2, 2026

View reviewed changes

claude Bot reviewed Aug 2, 2026

View reviewed changes

VibhuJawa added 2 commits August 2, 2026 01:20

refactor(interleaved): clarify legacy source migration

bbd65dd

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

refactor(schema): share Parquet FILE type

db36956

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Align interleaved source locators with Parquet FILE - #2284

Align interleaved source locators with Parquet FILE#2284
VibhuJawa wants to merge 7 commits into
NVIDIA-NeMo:mainfrom
VibhuJawa:codex/interleaved-file

VibhuJawa commented Aug 1, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 1, 2026

Uh oh!

greptile-apps Bot commented Aug 1, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VibhuJawa commented Aug 1, 2026

Uh oh!

VibhuJawa commented Aug 2, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

claude Bot Aug 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

VibhuJawa commented Aug 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Compatibility and scope

Performance validation

Validation

Diff

Uh oh!

copy-pr-bot Bot commented Aug 1, 2026

Uh oh!

greptile-apps Bot commented Aug 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VibhuJawa commented Aug 1, 2026

Uh oh!

VibhuJawa commented Aug 2, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

claude Bot Aug 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

VibhuJawa commented Aug 1, 2026 •

edited

Loading

greptile-apps Bot commented Aug 1, 2026 •

edited

Loading